from typing import Dict, Any, Optional
from src.core.output_manager_base import OutputManagerBase, SimulationParameters, MPI, fem, FENICSX_AVAILABLE
import numpy as np
import json
import os

class SolutionStatisticsCalculator(OutputManagerBase):
    """
    用于计算解的统计数据（例如，均值、最大值、最小值、范数等）。
    """
    def __init__(self, 
                 comm: Optional[MPI.Comm], 
                 params: SimulationParameters, 
                 fem_entities: Dict[str, Any], 
                 output_dir: str,
                 stats_frequency: int = 1,
                 functions_to_analyze: Optional[list[str]] = None):
        super().__init__(comm, params, fem_entities, output_dir)
        self.stats_frequency = stats_frequency
        self.functions_to_analyze_names = functions_to_analyze
        self.all_statistics: Dict[str, Dict[str, Any]] = {}
        
        if self.comm is None or self.comm.rank == 0: # type: ignore
            print(f"SolutionStatisticsCalculator: 初始化。统计频率: {self.stats_frequency}。")
            if self.functions_to_analyze_names:
                print(f"  将专门分析以下函数: {self.functions_to_analyze_names}")
            else:
                print("  将尝试分析 fem_entities 中的所有 fem.Function 对象。")

    def is_output_step(self, iteration_count: int, current_time: float) -> bool:
        return iteration_count % self.stats_frequency == 0

    def execute_step(self, iteration_count: int, current_time: float) -> None:
        if self.comm is None or self.comm.rank == 0: # type: ignore
            print(f"SolutionStatisticsCalculator: 在迭代 {iteration_count}, 时间 {current_time:.2f} 执行统计计算。")
        
        if not FENICSX_AVAILABLE:
            if self.comm is None or self.comm.rank == 0: # type: ignore
                print("  SolutionStatisticsCalculator: FEniCSx 不可用，跳过统计计算。")
            # 即使跳过，也应该记录迭代和时间
            self.all_statistics[f"iter_{iteration_count}"] = {"time": current_time}
            return

        current_iter_stats: Dict[str, Any] = {} 

        functions_to_process: Dict[str, fem.Function] = {} # type: ignore
        if self.functions_to_analyze_names:
            for name in self.functions_to_analyze_names:
                func = self.fem_entities.get(name)
                if isinstance(func, fem.Function): # type: ignore
                    functions_to_process[name] = func # type: ignore
                elif self.comm is None or self.comm.rank == 0: # type: ignore
                    print(f"  警告: 指定的函数 '{name}' 在 fem_entities 中未找到或不是 fem.Function 类型。")
        else:
            for key, entity in self.fem_entities.items():
                if isinstance(entity, fem.Function): # type: ignore
                    functions_to_process[key] = entity # type: ignore
        
        if not functions_to_process:
            if self.comm is None or self.comm.rank == 0: # type: ignore
                print("  SolutionStatisticsCalculator: 未找到可供分析的 fem.Function。")
            self.all_statistics[f"iter_{iteration_count}"] = {"time": current_time} 
            return

        for name, func in functions_to_process.items():
            func_stats: Dict[str, Any] = {}
            try:
                local_values = func.x.array
                
                func_stats["min_local"] = np.min(local_values) if local_values.size > 0 else None
                func_stats["max_local"] = np.max(local_values) if local_values.size > 0 else None
                func_stats["mean_local"] = np.mean(local_values) if local_values.size > 0 else None
                func_stats["std_local"] = np.std(local_values) if local_values.size > 0 else None
                
                if self.comm is not None and FENICSX_AVAILABLE:
                    min_val_to_reduce = func_stats["min_local"] if func_stats["min_local"] is not None else float('inf')
                    max_val_to_reduce = func_stats["max_local"] if func_stats["max_local"] is not None else float('-inf')

                    global_min = self.comm.allreduce(min_val_to_reduce, op=MPI.MIN) # type: ignore
                    global_max = self.comm.allreduce(max_val_to_reduce, op=MPI.MAX) # type: ignore
                    
                    from dolfinx.fem import form, assemble_scalar # type: ignore
                    from ufl import inner, dx # type: ignore
                    
                    mesh = self.fem_entities.get("mesh")
                    if mesh is None:
                        # 如果没有mesh，就不能计算全局L2范数
                        if self.comm.rank == 0:
                             print(f"  警告: 函数 '{name}' 无法计算全局L2范数，因为 fem_entities 中缺少 'mesh'。")
                        global_L2_norm = None
                    else:
                        norm_L2_form = form(inner(func, func) * dx(mesh)) # type: ignore
                        global_L2_norm_sq = assemble_scalar(norm_L2_form) # type: ignore
                    global_L2_norm = np.sqrt(self.comm.allreduce(global_L2_norm_sq, op=MPI.SUM)) # type: ignore
                    
                    func_stats["min_global"] = global_min if global_min != float('inf') else None
                    func_stats["max_global"] = global_max if global_max != float('-inf') else None
                    func_stats["norm_L2_global"] = global_L2_norm

                    if self.comm.rank == 0:
                        min_g_str = f"{global_min:.4e}" if global_min != float('inf') else "N/A"
                        max_g_str = f"{global_max:.4e}" if global_max != float('-inf') else "N/A"
                        l2_g_str = f"{global_L2_norm:.4e}" if global_L2_norm is not None else "N/A"
                        print(f"    函数 '{name}': Min(g)={min_g_str}, Max(g)={max_g_str}, L2Norm(g)={l2_g_str}")
                else:
                    if self.comm is None or self.comm.rank == 0: # type: ignore
                        min_local_str = f"{func_stats['min_local']:.4e}" if func_stats['min_local'] is not None else "N/A"
                        max_local_str = f"{func_stats['max_local']:.4e}" if func_stats['max_local'] is not None else "N/A"
                        print(f"    函数 '{name}' (局部统计): Min={min_local_str}, Max={max_local_str}")
                
                current_iter_stats[name] = func_stats
            except Exception as e:
                if self.comm is None or self.comm.rank == 0: # type: ignore
                    print(f"  错误: 计算函数 '{name}' 的统计数据时出错: {e}")
        
        current_iter_stats["time"] = current_time 
        self.all_statistics[f"iter_{iteration_count}"] = current_iter_stats

    def finalize(self) -> None:
        if self.comm is None or self.comm.rank == 0: # type: ignore
            print("SolutionStatisticsCalculator: 完成所有统计计算。")
            # 文件保存逻辑 (可选)
            stats_filename = getattr(self.params, 'stats_output_filename', None)
            if stats_filename:
                stats_file_path = os.path.join(self.output_dir, stats_filename)
                os.makedirs(self.output_dir, exist_ok=True)
                try:
                    with open(stats_file_path, 'w') as f:
                        json.dump(self.all_statistics, f, indent=4)
                    print(f"  所有统计数据已保存到: {stats_file_path}")
                except Exception as e:
                    print(f"  错误: 保存统计数据到 JSON 文件时出错: {e}")

            if self.all_statistics:
                print("  收集到的统计数据摘要:")
                for iter_key, iter_data_dict in self.all_statistics.items():
                    print(f"    {iter_key}:")
                    time_val = iter_data_dict.get("time")
                    if time_val is not None:
                        # 确保 time_val 是数字以便格式化
                        try:
                            print(f"      time: {float(time_val):.4e}") 
                        except (ValueError, TypeError):
                             print(f"      time: {time_val}") # 如果无法格式化，直接打印

                    for func_name, stats_data in iter_data_dict.items():
                        if func_name == "time": 
                            continue
                        if isinstance(stats_data, dict):
                            min_val = stats_data.get('min_global', stats_data.get('min_local'))
                            max_val = stats_data.get('max_global', stats_data.get('max_local'))
                            
                            min_val_str = f"{min_val:.4e}" if isinstance(min_val, (int, float)) else "N/A"
                            max_val_str = f"{max_val:.4e}" if isinstance(max_val, (int, float)) else "N/A"
                            
                            l2_norm_str = ""
                            l2_norm_val = stats_data.get('norm_L2_global')
                            if isinstance(l2_norm_val, (int, float)):
                                l2_norm_str = f", L2Norm(g)={l2_norm_val:.4e}"
                            
                            print(f"      {func_name}: Min={min_val_str}, Max={max_val_str}{l2_norm_str}")
                        else:
                            print(f"      {func_name}: {stats_data}") # 其他类型数据